import atexit
import os
import threading
import time
from tkinter import Frame

os.environ["MUJOCO_GL"] = "egl"
import mediapy as media
import mujoco.viewer
import numpy as np
from scipy.spatial.transform import Rotation

from sim.mj_ik import MujocoIK


class MujocoPiperSim(MujocoIK):
    """
    A simulator for the Piper robot.
    The movement is given by ik solver.

    The planner only needs to predict the xquat trajectories
    """

    def __init__(self, radius=0.04, reg=0.01, render=True):
        super().__init__(radius, reg)

        self.model.vis.scale.framewidth = 0.02  # linewidth of the coordinate axis
        # set camera
        self.camera = mujoco.MjvCamera()
        mujoco.mjv_defaultFreeCamera(self.model, self.camera)
        self.camera.distance = 1.5
        self.camera.elevation = -15
        self.camera.azimuth = -150
        self.camera.lookat = (0, 0, 0)

        # Visualize site frames and labels
        self.renderer = None
        self.voption = None
        self.init_render()

        # function to close the renderer when the program exits
        atexit.register(self._cleanup_on_exit)

        # Reset the state to the "zero_position" keyframe.
        self.go_zero()

    def _cleanup_on_exit(self):
        """
        Close the renderer when the program exits
        """
        if self.renderer is not None:
            self.renderer.close()
            self.renderer = None

    def init_render(self, height=540, width=960):

        self.renderer = mujoco.Renderer(self.model, height=height, width=width)
        self.voption = mujoco.MjvOption()
        # self.voption.frame = mujoco.mjtFrame.mjFRAME_SITE
        self.voption.label = mujoco.mjtLabel.mjLABEL_SITE
        self.renderer.update_scene(self.data, self.camera, self.voption)
        self.voption.label = mujoco.mjtLabel.mjLABEL_NONE

    def move_target_xeuler(self, pos: np.ndarray, rpy: np.ndarray):
        assert len(pos) == len(rpy) == 3, "pos and rpy should be a 3-dim array"
        mocap_id = self.model.body("target").mocapid
        self.data.mocap_pos[mocap_id] = pos
        quat = Rotation.from_euler("xyz", rpy, degrees=False).as_quat(scalar_first=True)
        self.data.mocap_quat[mocap_id] = quat
        mujoco.mj_kinematics(self.model, self.data)

    def get_effector_xquat(self):
        effector_quat = np.empty(4)
        mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)
        xquat = np.concat([self.data.site("effector").xpos, effector_quat])
        return xquat

    def get_target_xquat(self):
        target_quat = np.empty(4)
        mujoco.mju_mat2Quat(target_quat, self.data.body("target").xmat)
        xquat = np.concat([self.data.site("target").xpos, target_quat])
        return xquat

    def get_target_xeuler(self):
        target_euler = Rotation.from_quat(self.data.body("target").xquat, scalar_first=True).as_euler("xyz")
        return np.concat([self.data.site("target").xpos, target_euler])

    def get_effector_xeuler(self):
        effector_euler = Rotation.from_quat(self.data.body("effector").xquat, scalar_first=True).as_euler("xyz")
        return np.concat([self.data.body("effector").xpos, effector_euler])

    def et_error(self, joint=None, pos=None, rpy=None):
        """
        Effector-target error: perform forward kinematics and check the distance
        between the effector and the target.

        Currently only check the position distance.
        The orientation error is not considered yet
        """
        if joint is None:
            joint = self.data.qpos[:6]
        if pos is None:
            pos = self.data.site("target").xpos.copy()
        if rpy is None:
            rpy = Rotation.from_quat(self.data.body("target").xquat, scalar_first=True).as_euler("xyz")

        # local step without changing the data inplace
        data = mujoco.MjData(self.model)
        data.qpos = np.concat((joint, np.zeros(2)))
        mujoco.mj_kinematics(self.model, data)

        return np.linalg.norm(data.body("effector").xpos - self.get_target_xeuler()[:3])

    def go_home(self):
        key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, "home")
        mujoco.mj_resetDataKeyframe(self.model, self.data, key)
        mujoco.mj_forward(self.model, self.data)

    def go_zero(self):
        key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, "zero_position")
        mujoco.mj_resetDataKeyframe(self.model, self.data, key)
        mujoco.mj_forward(self.model, self.data)

    def close_gripper(self):
        joints = np.concat([self.data.qpos[:6], np.zeros(2)])
        return self.step(joints)

    def open_gripper(self):
        joints = np.concat([self.data.qpos[:6], np.array([1, -1]) * 0.035])
        return self.step(joints)

    def render_frame(self):
        if self.renderer is None:
            print("Warning: Renderer not initialized. Call init_render() first.")
            return None
        self.renderer.update_scene(self.data, self.camera, self.voption)
        return self.renderer.render()

    def render(self):
        self.step()
        if self.renderer is not None:
            media.show_image(self.render_frame())
        else:
            print("Warning: Renderer not initialized. Call init_render() first.")

    def mujoco_viewer(self, run_in_background=True):
        """
        Used for notebook running
        """

        def _viewer_loop():
            with mujoco.viewer.launch_passive(self.model, self.data) as viewer:
                print("Viewer launched!")
                while viewer.is_running():
                    self.step()
                    # mujoco.mj_step(self.model, self.data)
                    viewer.sync()
                    time.sleep(0.01)
                print("Viewer closed!")

        if run_in_background:
            # daemon=True: thread will be terminated when the main thread terminates
            self.viewer_thread = threading.Thread(target=_viewer_loop, daemon=True)
            self.viewer_thread.start()
            print("Viewer started in background thread. You can now run other cells!")
            return self.viewer_thread
        else:
            _viewer_loop()

    def step(self, x: np.ndarray = None):
        """
        x: the joint angles to render.
        If x is None, render the current status of the model
        It first update the kinematics of the model, then update the scene,
        and then render the scene.
        image render: media.show_image(self.render())
        """
        if x is not None:
            self.data.qpos = x
        mujoco.mj_kinematics(self.model, self.data)
        mujoco.mj_camlight(self.model, self.data)

    def move_to(self, pos: np.ndarray, rpy: np.ndarray):
        """
        reset the target position self.data.site("target") to (pos, quat)
        rpy: row pitch yaw in radians
        move to the position using the ik solver
        """

        # self.move_target_xeuler(pos, rpy)

        assert len(pos) == len(rpy) == 3, "pos and rpy should be a 3-dim array"
        j0 = self.data.qpos.copy()
        joints = self.solve_ik_xyz_rpy(xyz=pos, rpy=rpy, j0=j0, x_prev=j0.copy())
        self.step(joints)

    def planning_to(self, planning_model, pos: np.ndarray, rpy: np.ndarray, render: bool = False):

        self.move_target_xeuler(pos, rpy)

        cond = np.concat([self.get_effector_xeuler(), pos, rpy], axis=-1)
        # (batch_size, num_act_samples, horizon_steps, xquat_dim=7)
        path = planning_model(cond)[0, 0]
        if render and self.renderer is not None:
            frames = [self.render_frame()]
        else:
            frames = []
        for pt in path:
            self.move_to(pt[:3], pt[3:])
            if render and self.renderer is not None:
                frames.append(self.render_frame())
        return frames


if __name__ == "__main__":

    def test_simulator():
        self = MujocoPiperSim()
        pos = np.array([0.34803, 0.21306, 0.0])
        rpy = np.array([0, 1.7, 0])
        self.move_target_xeuler(pos, rpy)
        self.move_to(pos=pos, rpy=rpy)
        self.render()
        self.go_home()

    def check_jacobian():
        """
        the un-change of other parameters (say, rpy) may cause the problem.
        """
        self = MujocoPiperSim()
        nq = self.model.nq  # number of joints

        self.go_zero()

        # Disable gravity
        self.model.opt.gravity = np.zeros(3)
        # lower damping
        self.model.dof_damping = self.model.dof_damping * 0.1

        frames = []
        for i in range(10000):
            jacp = np.empty((3, nq))  # position xyz
            jacr = np.empty((3, nq))  # Lie algebra w1, w2, w3

            mujoco.mj_forward(self.model, self.data)
            mujoco.mj_comPos(self.model, self.data)
            mujoco.mj_jacSite(self.model, self.data, jacp, jacr, self.data.site("effector").id)

            jac = np.vstack([jacp, jacr])[:, :6]

            jac_inv = np.linalg.inv(jac.T @ jac + 0.1 * np.eye(nq - 2)) @ jac.T

            dx = np.array([0.001, 0, 0, 0, 0, 0])

            dq = jac_inv @ dx

            # self.data.ctrl = self.data.qpos + np.concat([dq, np.zeros(2)])
            # mujoco.mj_step(self.model, self.data)

            self.data.qpos = self.data.qpos + np.concat([dq, np.zeros(2)])
            self.data.qpos = np.clip(
                self.data.qpos, self.model.actuator_ctrlrange[:, 0], self.model.actuator_ctrlrange[:, 1]
            )

            mujoco.mj_kinematics(self.model, self.data)
            mujoco.mj_comPos(self.model, self.data)

            if i % 10 == 0:
                frames.append(self.render_frame())

        media.show_video(frames, loop=False)

        # dq = np.linalg.solve(jacp, dx)

    def check_data():
        from data.data_loaders import TrajectoryPlanningData

        self = MujocoPiperSim()
        data = TrajectoryPlanningData("data/json_data/sim_data_1M_0.1.json", horizon_steps=20)
        frames = []
        for xeuler in data[0].action:
            self.move_to(xeuler[:3], xeuler[3:])
            self.move_target_xeuler(xeuler[:3], xeuler[3:])
            frames.append(self.render_frame())

        media.show_video(frames, loop=False)

    def test_model():
        import hydra
        from omegaconf import DictConfig

        with hydra.initialize(version_base=None, config_path="cfg"):
            ej_cfg: DictConfig = hydra.compose(
                config_name="train_diffusion_planning",
                overrides=[
                    "tag=fixed",
                    "model.temperature=0.",
                    "horizon_steps=20",
                    "model.scale=1",
                ],
            )
            model = hydra.utils.instantiate(ej_cfg.model)
            model.load_ckpt(ej_cfg.ckp_name + ".ckp")
            model.num_act_samples = 1
        self = MujocoPiperSim()

        pos = np.array([0.54803, 0.27306, 0.0])
        # (w, x, y, z)
        rpy = np.array([0, 2, 0])

        # from true data
        # pos = np.array([0.5973, -0.4459, -0.0455])
        # quat = np.array([0.5167, 0.1118, 0.8023, -0.2771])

        # planning
        self.go_zero()
        frames = [self.render_frame()]
        for i in range(1):
            frames += self.planning_to(model, pos=pos, rpy=rpy, render=True)
        media.show_video(frames, loop=False)
